import os
import torch.utils.data
import cv2
import numpy as np
import imgaug as ia
from imgaug import augmenters as iaa
from libs.utils import getFilepath, np2Tensor
from libs.utils import show_sample_withplt, show_sample_withocv


class HairDatasetAug(torch.utils.data.Dataset):
    def __init__(self, data_folder, imglist, image_size=(480, 480), mask_size=(480, 480), blurmask_size=(480, 480),
                 trimap_size=(480, 480), alpha_size=(480, 480), return_blurmask=False, return_mask=True,
                 return_trimap=False, return_alpha=True):
        self.data_folder = data_folder
        if not os.path.exists(self.data_folder):
            raise Exception("%s  not exists." % self.data_folder)

        self.image_list, self.blurmask_list, self.mask_list, self.trimap_list, self.alpha_list = getFilepath(data_folder, imglist)
        self.image_size = image_size
        self.mask_size = mask_size
        self.blurmask_size = blurmask_size
        self.trimap_size = trimap_size
        self.alpha_size = alpha_size
        self.return_blurmask = return_blurmask
        self.return_mask = return_mask
        self.return_trimap = return_trimap
        self.return_alpha = return_alpha

    def activator_binmasks(images, image, augmenter, parents, default):
        # forbid these augumenters for mask images
        if augmenter.name in ["Multiply", "GaussianBlur", "AverageBlur", "MotionBlur"]:
            return False
        else:
            return default

    def image_aug(self, image, blurmask, mask, trimap, alpha):
        blurmask_aug = None
        mask_aug = None
        trimap_aug = None
        alpha_aug = None
        hooks_binmasks = ia.HooksImages(activator=self.activator_binmasks)

        seq = iaa.SomeOf((1, 4),
                         [   iaa.Multiply((0.4, 1.5), name="Multiply"),
                             iaa.Fliplr(0.5, name="Flip"),
                             iaa.SomeOf((1, 3), [iaa.Affine(translate_percent={"x": (-0.2, 0.2), "y": (-0.2, 0.2)}, cval=0, name="translate_percent"),
                                                 iaa.Affine(rotate=(-15, 15), cval=0, name="Rotate"),
                                                 iaa.Affine(scale=(0.6, 1.3), cval=0, name="Scale"),
                                                 iaa.Affine(shear=(-15, 15), cval=0, name="Sheer")]),
                             iaa.OneOf([iaa.GaussianBlur((0, 3.0), name="GaussianBlur"),
                                        iaa.AverageBlur(k=(0, 3), name="AverageBlur"),
                                        iaa.MotionBlur(k=6, name="MotionBlur")
                                        ])
                         ])
        # call this for each batch again, NOT only once at the start to ensure same augumentation applied to image and mask
        seq_det = seq.to_deterministic()
        image_aug = np.squeeze(seq_det.augment_images(np.expand_dims(np.array(image), axis=0)), axis=0)
        if self.return_blurmask:
            blurmask_aug = np.squeeze(seq_det.augment_images((np.expand_dims(np.array(blurmask), axis=0)), hooks=hooks_binmasks), axis=0)
        if self.return_mask:
            mask_aug = np.squeeze(seq_det.augment_images((np.expand_dims(np.array(mask), axis=0)), hooks=hooks_binmasks), axis=0)
        if self.return_trimap:
            trimap_aug = np.squeeze(seq_det.augment_images((np.expand_dims(np.array(trimap), axis=0)), hooks=hooks_binmasks), axis=0)
        if self.return_alpha:
            alpha_aug = np.squeeze(seq_det.augment_images((np.expand_dims(np.array(alpha), axis=0)), hooks=hooks_binmasks), axis=0)

        return image_aug, blurmask_aug, mask_aug, trimap_aug, alpha_aug

    def image_preprocess(self, image):
        image = cv2.resize(image, self.image_size)
        image = image.astype(np.float32)
        image = ((image / 255.0) - 0.5) / 0.5
        image = np2Tensor(image)
        return image

    def blurmask_preprocess(self, blurmask):
        blurmask = cv2.resize(blurmask, self.blurmask_size)
        blurmask = blurmask.astype(np.float32)
        blurmask = (blurmask / 255.0)
        blurmask = np2Tensor(blurmask)
        return blurmask

    def mask_preprocess(self, mask):
        mask = cv2.resize(mask, self.mask_size)
        mask = mask.astype(np.float32)
        mask = (mask / 255.0)
        mask = np2Tensor(mask)
        return mask

    def trimap_preprocess(self, trimap):
        trimap = cv2.resize(trimap, self.trimap_size)
        trimap[trimap == 0] = 0
        trimap[trimap == 128] = 1
        trimap[trimap == 255] = 2
        trimap = np2Tensor(trimap)
        return trimap

    def alpha_preprocess(self, alpha):
        alpha = cv2.resize(alpha, self.alpha_size)
        alpha = alpha.astype(np.float32)
        alpha = (alpha / 255.0)
        alpha = np2Tensor(alpha)
        return alpha

    def __getitem__(self, index):
        output_dict = {}

        image, blurmask, mask, trimap, alpha = self.read_all_image_files(index)
        image, blurmask, mask, trimap, alpha = self.image_aug(image, blurmask, mask, trimap, alpha)

        output_dict['image'] = self.image_preprocess(image)

        if self.return_blurmask:
            output_dict['blurmask_gt'] = self.blurmask_preprocess(blurmask)

        if self.return_mask:
            output_dict['mask_gt'] = self.mask_preprocess(mask)

        if self.return_trimap:
            output_dict['trimap_gt'] = self.trimap_preprocess(trimap)

        if self.return_alpha:
            output_dict['alpha_gt'] = self.alpha_preprocess(alpha)

        return output_dict

    def __len__(self):
        return len(self.image_list)

    def read_all_image_files(self, index):

        blurmask = None
        mask = None
        trimap = None
        alpha = None

        img_path = self.image_list[index]
        image = cv2.imread(img_path)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        if self.return_blurmask:
            blurmask_path = self.blurmask_list[index]
            blurmask = cv2.imread(blurmask_path, 0)

        if self.return_mask:
            mask_path = self.mask_list[index]
            mask = cv2.imread(mask_path, 0)

        if self.return_trimap:
            trimap_path = self.trimap_list[index]
            trimap = cv2.imread(trimap_path, 0)

        if self.return_alpha:
            alpha_path = self.alpha_list[index]
            alpha = cv2.imread(alpha_path, 0)

        return image, blurmask, mask, trimap, alpha



if __name__ == "__main__":

    return_blurmask = True
    return_mask = True
    return_trimap = True
    return_alpha = True
    input_shape = (480, 480)
    output_shape = (480, 480)
    for sample in HairDatasetAug('/home/hanson/work/HairMattingTools/train',
                                '/home/hanson/work/HairMattingTools/testlist.txt',
                                 image_size=input_shape,
                                 mask_size=output_shape,
                                 blurmask_size=output_shape,
                                 trimap_size=output_shape,
                                 alpha_size=output_shape, return_blurmask=return_blurmask,
                                 return_mask=return_mask, return_trimap=return_trimap, return_alpha=return_alpha):
        show_sample_withocv(sample)